

class AVLTreeNode_BF:
    SKEW_LEFT = 1       # 左子树比右子树高度多1
    SKEW_BAL = 0        # 左右平衡
    SKEW_RIGHT = -1     # 与1相反

    DEPTH_PLUS = 1
    DEPTH_STAY = 0

    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
        self.balance_factor = self.SKEW_BAL     # 平衡因子: 左子树高度 - 右子树高度

    def insert(self, val):
        # 返回旋转后新的节点, 和树增加的高度
        if val == self.val:
            return self, self.DEPTH_STAY

        if val < self.val:
            if self.left:
                new_left, depth_plus = self.left.insert(val)
                return self.bal_left(new_left, depth_plus, val)
            else:
                # 创建新节点
                self.left = AVLTreeNode(val)
                if self.right:
                    self.balance_factor = 0
                    return self, self.DEPTH_STAY
                else:
                    self.balance_factor = 1
                    return self, self.DEPTH_PLUS
        else:
            # 和self.left对称
            # print('here')
            assert val > self.val
            if self.right:
                new_right, depth_plus = self.right.insert(val)
                return self.bal_right(new_right, depth_plus, val)
            else:
                # 添加self.right节点
                # balance_factor根据是否具有self.left节点决定
                self.right = AVLTreeNode(val)
                if self.left:
                    self.balance_factor = 0
                    return self, self.DEPTH_STAY
                else:
                    self.balance_factor = -1
                    return self, self.DEPTH_PLUS

    def bal_left(self, new_left, depth_plus, val):
        # 根据子树插入节点后,返回的增加高度, 计算当前节点的平衡因子.
        # 根据平衡因子判断是否需要旋转, 并根据val判断如何进行旋转.
        old_left = self.left
        old_bf = self.balance_factor
        self.balance_factor += depth_plus       # 如果调整后, 这里还要继续变化
        new_bf = self.balance_factor

        # depth_plus in [0, 1]: 左子树高度增加或者不变(导致当前节点平衡因子增加或者不变)
        # 左边插入: bf变化可能性有6种: -1 到0/1, 0到0/1, 1到1/2
        # 如果 bf不变, 表示当前节点增加高度为0
        # 如果 新bf为0, 表示当前节点增加高度为0(填补了子树中一些层次不齐的地方而已)
        # 如果 新bf为2, 需要旋转
        # 如果 old_bf为0, 新bf为1, 表示当前节点高度增加1

        self.left = new_left
        if new_bf <= 1:
           # 不需要旋转
            # assert new_bf in [0, 1]
            if old_bf == 0 and new_bf == 1:
                return self, 1
            else:
                assert old_bf == new_bf or new_bf == 0
                return self, 0
        else:
            # 需要旋转(单旋或双旋)
            assert new_bf == 2
            assert self.left is old_left
            return self.rotate_left(val)

    def rotate_left(self, val):
        # 分为单旋和双旋
        if val < self.left.val:
            return self._rotate_left_left(val)
        else:
            return self._rotate_left_right(val)

    def _rotate_left_left(self, val):
        '''
        单旋
        k2 > k1
        k1 在没有插入val前, k1的深度比k2右子树深1
        k1 在插入val后(插入到了k1的左子树), k1的深度比k2由子树深2(导致k2, 也就是当前节点的平衡因子变为2)
        1: 所以(k1 + k1的右子树)深度比k2的右子树深1
        2: 在插入val后, 所以k1的左子树深度比k1的右子树深1(k1+k1的左子树比k2的右子树深2)
        k1的右子树(不包括k1)高度和k2的右子树高度相同.
        插入val后k1的左子树(不包括k1)高度比k1的右子树高度高1
        现在调整后的k2/self 左右包括(k1.right, k2.right): 所以旋转后的k2的平衡因子为0
        现在调整后的k1高度比k2提升了1(k1成为了新的root), k1的左右包括(k1.left, k2). k1上升的1高度, 抵消了k2子树比k1左子树低的1
        所以旋转后的k1平衡因子为0
        '''
        k2 = self
        k1 = self.left

        print('single left k1, k2 [%s, %s]' % (k1.val, k2.val))

        save_k1_right = k1.right
        k1.right = k2
        k2.left = save_k1_right

        # 重新设置 k1, k2 balance_factor(说明见注释: 或者参阅avl图示)
        k1.balance_factor = 0
        k2.balance_factor = 0

        return k1, 0

    def _rotate_left_right(self, val):
        # k3 > k2 > k1
        print('rotate left right')
        assert val > self.left.val
        k3 = self
        k1 = self.left
        k2 = self.left.right

        # print(k3, k1, k2)

        # rotate
        save_k2_left = k2.left
        save_k2_right = k2.right

        k2.left = k1
        k2.right = k3

        k1.right = save_k2_left
        k3.left = save_k2_right

        # adject balance factor
        save_k2_bf = k2.balance_factor
        k2.balance_factor = 0

        assert save_k2_bf in [-1, 0, 1]
        if save_k2_bf == 1:
            k1.balance_factor = 0
            k3.balance_factor = -1
        elif save_k2_bf == -1:
            k1.balance_factor = 1
            k3.balance_factor = 0
        else:
            # 如果k2.bf == 0, 说明k2是新插入的节点
            assert save_k2_left is None and save_k2_right is None
            k1.balance_factor = 0
            k3.balance_factor = 0

        return k2, 0

    def bal_right(self, new_right, depth_plus, val):
        old_right = self.right

        # 因为是右边高度增长,
        # 所以用- depth_plus 作为新的balance_factor
        old_bf = self.balance_factor
        self.balance_factor -= depth_plus
        new_bf = self.balance_factor
        self.right = new_right
        if new_bf >= -1:
            # assert new_bf in [-1, 0]
            if new_bf == old_bf or new_bf == 0:
                return self, 0
            else:
                assert old_bf == 0 and new_bf == -1
                return self, 1
        else:
            print('new_bf %s' % new_bf)
            assert new_bf == -2
            assert self.right is old_right
            return self.rotate_right(val)

    def rotate_right(self, val):
        if val > self.right.val:
            return self._rotate_right_right(val)
        else:
            return self._rotate_right_left(val)

    def _rotate_right_right(self, val):
        # 单旋
        # k1 < k2
        print('===================  =====================')
        print('single right')
        print('self %s' % self.val)
        print('self.right %s' % self.right.val)
        print('self.rr %s' % self.right.right.val)
        k1 = self
        k2 = self.right

        # print('k1 %s' % k1.val)

        save_k2_left = k2.left
        k2.left = k1
        k1.right = save_k2_left

        # 重新设置 k1, k2 balance_factor
        k1.balance_factor = 0
        k2.balance_factor = 0

        # 并且调整后的k2所代表的树高, 相对于插入前k1/self所代表的树高并没有增长. 所以返回depth_plus = 0
        return k2, 0

    def _rotate_right_left(self, val):
        # k3 > k2 > k1
        print('rotate right left')
        assert val < self.right.val
        k1 = self
        k3 = self.right
        k2 = self.right.left

        print('k1, k2, k3 [%s, %s, %s]' % (k1.val, k2.val, k3.val))

        # rotate
        save_k2_left = k2.left
        save_k2_right = k2.right

        k2.left = k1
        k2.right = k3

        k1.right = save_k2_left
        k3.left = save_k2_right

        # adject balance factor
        save_k2_bf = k2.balance_factor
        k2.balance_factor = 0

        assert save_k2_bf in [-1, 0, 1]        # k2 一定不平衡(如果平衡: 那么k2就不会造成self/k1的不平衡)
        if save_k2_bf == 1:
            k1.balance_factor = 0
            k3.balance_factor = -1
        elif save_k2_bf == -1:
            k1.balance_factor = 1
            k3.balance_factor = 0
        else:
            assert save_k2_left is None and save_k2_right is None
            k1.balance_factor = 0
            k3.balance_factor = 0

        # 并且调整后的k2所代表的树高, 相对于插入前k1/self所代表的树高并没有增长. 所以返回depth_plus = 0
        return k2, 0

    def rec_show(self, cnt):
        if self.left:
            cnt = self.left.rec_show(cnt)

        print('%s:%s' % (self.val, self.balance_factor), end=' ', sep=' ')
        cnt += 1

        if self.right:
            cnt = self.right.rec_show(cnt)

        return cnt


class AVLTree:
    def __init__(self):
        self.root = None

    def insert(self, val):
        if self.root is None:
            self.root = AVLTreeNode(val)
            return self.root

        else:
            self.root, _  = self.root.insert(val)

    def show(self):
        if self.root:
            cnt = self.root.rec_show(cnt=0)
        print('\n=================== cnt %s =====================' % cnt)
